#!/usr/bin/env python3
"""
linux-sgx MCP Server

面向软件：linux-sgx（Intel Software Guard Extensions）
所有 MCP tool 列表：

1. aesmd_service_control
   - 用途：管理 aesmd 服务的启动、停止、重启和重载操作
   - 参数：action (str) - 操作类型：start/stop/restart/reload/status
   - 返回结构：统一 JSON 结构

2. remount_dev_exec_service_control
   - 用途：管理 remount-dev-exec 服务的启动、停止、重启和重载操作
   - 参数：action (str) - 操作类型：start/stop/restart/reload/status
   - 返回结构：统一 JSON 结构
"""

import subprocess
import json
from typing import Dict, Any
from mcp.server.fastmcp import FastMCP, Context

mcp = FastMCP("Linux SGX MCP Server")

@mcp.tool()
def aesmd_service_control(action: str) -> Dict[str, Any]:
    """
    管理 aesmd 服务的各种操作
    
    参数:
        action: 操作类型，支持 start/stop/restart/reload/status
        
    返回 JSON 结构:
        {
            "success": bool,      # 操作是否成功
            "command": str,       # 执行的完整命令
            "exit_code": int,     # 命令退出码
            "stdout": str,        # 标准输出
            "stderr": str         # 标准错误输出
        }
    """
    valid_actions = ["start", "stop", "restart", "reload", "status"]
    if action not in valid_actions:
        return {
            "success": False,
            "command": f"systemctl {action} aesmd.service",
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Invalid action: {action}. Supported actions: {valid_actions}"
        }
    
    command = f"systemctl {action} aesmd.service"
    try:
        result = subprocess.run(
            command,
            shell=True,
            capture_output=True,
            text=True,
            timeout=30
        )
        
        return {
            "success": result.returncode == 0,
            "command": command,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command timed out after 30 seconds"
        }
    except Exception as e:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Command execution failed: {str(e)}"
        }

@mcp.tool()
def remount_dev_exec_service_control(action: str) -> Dict[str, Any]:
    """
    管理 remount-dev-exec 服务的各种操作
    
    参数:
        action: 操作类型，支持 start/stop/restart/reload/status
        
    返回 JSON 结构:
        {
            "success": bool,      # 操作是否成功
            "command": str,       # 执行的完整命令
            "exit_code": int,     # 命令退出码
            "stdout": str,        # 标准输出
            "stderr": str         # 标准错误输出
        }
    """
    valid_actions = ["start", "stop", "restart", "reload", "status"]
    if action not in valid_actions:
        return {
            "success": False,
            "command": f"systemctl {action} remount-dev-exec.service",
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Invalid action: {action}. Supported actions: {valid_actions}"
        }
    
    command = f"systemctl {action} remount-dev-exec.service"
    try:
        result = subprocess.run(
            command,
            shell=True,
            capture_output=True,
            text=True,
            timeout=30
        )
        
        return {
            "success": result.returncode == 0,
            "command": command,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command timed out after 30 seconds"
        }
    except Exception as e:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Command execution failed: {str(e)}"
        }

if __name__ == "__main__":
    mcp.run()